sleap_nn/config/training_job_config.py (3)

154-162: Missing docstring for new function

Unlike other functions in this file, load_sleap_config lacks a docstring describing its purpose, parameters, and return value.

Add a descriptive docstring:

def load_sleap_config(cls, json_file_path: str) -> TrainerConfig:
+    """Load a training job configuration from a legacy SLEAP JSON config file.
+
+    Arguments:
+        json_file_path: Path to a legacy SLEAP JSON configuration file.
+
+    Returns:
+        A TrainingJobConfig instance containing the mapped configuration data.
+    """
    with open(json_file_path, "r") as f:
        old_config = json.load(f)

---

158-161: Missing error handling for mapper functions

There's no error handling for potential exceptions in the mapper functions.

Add error handling for the mapper functions:

-    data_config = data_mapper(old_config)
-    model_config = model_mapper(old_config)
-    trainer_config = trainer_mapper(old_config)
+    try:
+        data_config = data_mapper(old_config)
+        model_config = model_mapper(old_config)
+        trainer_config = trainer_mapper(old_config)
+    except Exception as e:
+        raise ValueError(f"Failed to map configuration from {json_file_path}: {str(e)}")

---

162-162: Filename attribute not set in the created config

Unlike the YAML loading methods, this function doesn't set the filename attribute to the source file path.

Set the filename attribute in the returned config:

-    return cls(data_config=data_config, model_config=model_config, trainer_config=trainer_config)
+    return cls(
+        data_config=data_config, 
+        model_config=model_config, 
+        trainer_config=trainer_config,
+        filename=json_file_path
+    )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (1)

189-243: ⚠️ Potential issue

Inconsistency in parameter handling

The function implementation has several issues that need to be addressed:

1. Many parameters are commented out (lines 191-201) including critical path parameters like train_labels_path and val_labels_path which are marked as MISSING in the DataConfig class.
2. Line 210 comments out use_augmentations_train, but it's used in the conditional at line 241.
3. The structure of the legacy_config appears inconsistent with some values coming from data.preprocessing and others from optimization.augmentation_config.

def data_mapper(legacy_config: dict) -> DataConfig:
    return DataConfig(
-        # train_labels_path=legacy_config.get("train_labels_path", MISSING),
-        # val_labels_path=legacy_config.get("val_labels_path", MISSING),
+        train_labels_path=legacy_config.get("train_labels_path", MISSING),
+        val_labels_path=legacy_config.get("val_labels_path", MISSING),
        # test_file_path=legacy_config.get("test_file_path"),
        # provider=legacy_config.get("provider", "LabelsReader"),
        # user_instances_only=legacy_config.get("user_instances_only", True),
        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
        # np_chunks_path=legacy_config.get("np_chunks_path"),
        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
        # chunk_size=int(legacy_config.get("chunk_size", 100)),
        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),
        preprocessing=PreprocessingConfig(
            is_rgb=legacy_config.get("data", {}).get("preprocessing", {}).get("ensure_rgb", False),
            max_height=legacy_config.get("data", {}).get("preprocessing", {}).get("target_height"),
            max_width=legacy_config.get("data", {}).get("preprocessing", {}).get("target_width"),
            scale=legacy_config.get("data", {}).get("preprocessing", {}).get("input_scaling", 1.0),
            crop_hw=legacy_config.get("data", {}).get("preprocessing", {}).get("crop_size"),
            min_crop_size=legacy_config.get("data", {}).get("preprocessing", {}).get("crop_size_detection_padding", 100),
        ),
-        # use_augmentations_train=legacy_config.get("use_augmentations_train", False),
+        use_augmentations_train=legacy_config.get("use_augmentations_train", False),
        augmentation_config=AugmentationConfig(
            intensity=IntensityConfig(
                uniform_noise_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("uniform_noise_min_val", 0.0),

sleap_nn/config/data_config.py (1)

189-190: Add docstring to data_mapper function

The data_mapper function lacks a docstring to explain its purpose, parameters, and return value. This is important for maintainability and understanding the code's intent.

+"""
+Convert a legacy SLEAP configuration dictionary to a DataConfig object.
+
+Args:
+    legacy_config: Dictionary containing configuration from older SLEAP version
+
+Returns:
+    DataConfig: A structured configuration object matching the legacy settings
+"""
def data_mapper(legacy_config: dict) -> DataConfig:

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/model_config.py (1)

906-917: ⚠️ Potential issue

Always creating centroid head violates one-of semantics.

Because centroid is set unconditionally, it will be active alongside any other head, contradicting the @oneof policy that only one head can be non-None. Wrap this in a condition to ensure that we only create the centroid head if head_type is set to "centroid".

-            centroid=CentroidConfig(
+            centroid=(
+                CentroidConfig(
+                    confmaps=CentroidConfMapsConfig(
+                        anchor_part=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("anchor_part"),
+                        sigma=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("sigma", 5.0),
+                        output_stride=legacy_config.get("heads", {})
+                        .get("centroid", {})
+                        .get("output_stride", 1),
+                    )
+                )
+                if legacy_config.get("head_type") == "centroid"
+                else None
+            ),

sleap_nn/config/data_config.py (1)

270-277: ⚠️ Potential issue

Scale set to (None, None) may break validation logic.

By default, GeometricConfig.scale expects numeric tuples (e.g., (0.9, 1.1)). Using (None, None) can trigger type errors or produce unintended outcomes. Provide valid numeric values if you want to enable scaling, or set this to None if scaling is not used.

-   scale=(
-       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_min", None),
-       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_max", None),
-   ),
+   scale=(
+       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_min", 0.9),
+       legacy_config.get("optimization", {}).get("augmentation_config", {}).get("scale_max", 1.1),
+   ),

sleap_nn/config/model_config.py (1)

835-836: Remove or clarify leftover comment.

These lines appear to contain a partial or uncertain comment about pretrained_backbone_weights (e.g., )?? # i think its different). This can be confusing to future maintainers. Consider removing or clarifying these comments to ensure the code is clean and unambiguous.

-        # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")?? # i think its different
+        # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")  # Clarify if needed or remove entirely

sleap_nn/config/data_config.py (1)

223-224: use_augmentations_train is commented out, causing mismatch.

The DataConfig instance sets use_augmentations_train to its default False, even if legacy_config indicates otherwise. Uncomment and map use_augmentations_train from legacy_config to keep the final config in sync.

-        # use_augmentations_train=legacy_config.get("use_augmentations_train", False),
+        use_augmentations_train=legacy_config.get("use_augmentations_train", False),

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/trainer_config.py (6)

236-236: Add a docstring to the trainer_mapper function.

The function is missing documentation explaining its purpose, parameters, and return value. A docstring would improve maintainability and help other developers understand the function's role.

 def trainer_mapper(legacy_config: dict) -> TrainerConfig:
+    """Convert a legacy configuration dictionary to a TrainerConfig instance.
+
+    This function extracts values from a legacy SLEAP configuration dictionary
+    and maps them to the appropriate fields in the TrainerConfig structure.
+
+    Args:
+        legacy_config: A dictionary containing legacy configuration parameters.
+
+    Returns:
+        A TrainerConfig instance populated with values from the legacy config.
+    """
     return TrainerConfig(

---

243-246: Decide on commented-out code.

There are numerous commented-out sections throughout the function. These make the code harder to read and maintain. Either implement these sections if they're needed or remove them if they're not.

If these configurations are planned for future implementation, consider adding a TODO comment indicating this, or remove them entirely if they're not needed.

---

311-350: Simplify conditional configuration creation.

The conditional creation of lr_scheduler is complex and nested. Consider extracting this logic to a separate helper function to improve readability.

+        lr_scheduler=_create_lr_scheduler_from_legacy(legacy_config),
-        lr_scheduler=(
-            LRSchedulerConfig(
-                # step_lr=StepLRConfig(
-                #     step_size=legacy_config.get("optimization", {})
-                #     .get("lr_scheduler", {})
-                #     .get("step_lr", {})
-                #     .get("step_size", 10),
-                #     gamma=legacy_config.get("optimization", {})
-                #     .get("lr_scheduler", {})
-                #     .get("step_lr", {})
-                #     .get("gamma", 0.1),
-                # ) if legacy_config.get("optimization", {}).get("lr_scheduler", {}).get("scheduler") == "StepLR" else None,
-                reduce_lr_on_plateau=ReduceLROnPlateauConfig(
-                    # threshold=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("threshold", 1e-4),
-                    # threshold_mode=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("threshold_mode", "rel"),
-                    # cooldown=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("cooldown", 0),
-                    patience=legacy_config.get("optimization", {})
-                    .get("learning_rate_schedule", {})
-                    .get("plateau_patience", 10),
-                    # factor=legacy_config.get("optimization", {})
-                    # .get("lr_scheduler", {})
-                    # .get("reduce_lr_on_plateau", {})
-                    # .get("factor", 0.1),
-                    min_lr=legacy_config.get("optimization", {})
-                    .get("learning_rate_schedule", {})
-                    .get("min_learning_rate", 0.0),
-                )
-            )
-            if legacy_config.get("optimization", {}).get("learning_rate_schedule")
-            else None
-        ),

And add this helper function above:

def _create_lr_scheduler_from_legacy(legacy_config: dict) -> Optional[LRSchedulerConfig]:
    """Create LRSchedulerConfig from legacy configuration.
    
    Args:
        legacy_config: A dictionary containing legacy configuration parameters.
        
    Returns:
        A LRSchedulerConfig instance or None if learning_rate_schedule is not specified.
    """
    optimization = legacy_config.get("optimization", {})
    learning_rate_schedule = optimization.get("learning_rate_schedule")
    
    if not learning_rate_schedule:
        return None
        
    return LRSchedulerConfig(
        reduce_lr_on_plateau=ReduceLROnPlateauConfig(
            patience=learning_rate_schedule.get("plateau_patience", 10),
            min_lr=learning_rate_schedule.get("min_learning_rate", 0.0),
        )
    )

---

351-365: Similar to the lr_scheduler, extract early_stopping config creation to a helper function.

The conditional creation of early_stopping follows the same pattern as lr_scheduler and could benefit from the same refactoring approach.

+        early_stopping=_create_early_stopping_from_legacy(legacy_config),
-        early_stopping=(
-            EarlyStoppingConfig(
-                stop_training_on_plateau=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("reduce_on_plateau", False),
-                min_delta=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("plateau_min_delta", 0.0),
-                patience=legacy_config.get("optimization", {})
-                .get("learning_rate_schedule", {})
-                .get("plateau_patience", 1),
-            )
-            if legacy_config.get("optimization", {}).get("learning_rate_schedule")
-            else None
-        ),

And add this helper function below the _create_lr_scheduler_from_legacy function:

def _create_early_stopping_from_legacy(legacy_config: dict) -> Optional[EarlyStoppingConfig]:
    """Create EarlyStoppingConfig from legacy configuration.
    
    Args:
        legacy_config: A dictionary containing legacy configuration parameters.
        
    Returns:
        An EarlyStoppingConfig instance or None if learning_rate_schedule is not specified.
    """
    optimization = legacy_config.get("optimization", {})
    learning_rate_schedule = optimization.get("learning_rate_schedule")
    
    if not learning_rate_schedule:
        return None
        
    return EarlyStoppingConfig(
        stop_training_on_plateau=learning_rate_schedule.get("reduce_on_plateau", False),
        min_delta=learning_rate_schedule.get("plateau_min_delta", 0.0),
        patience=learning_rate_schedule.get("plateau_patience", 1),
    )

---

236-366: Add input validation for the legacy_config parameter.

The function doesn't validate the input legacy_config parameter. Consider adding validation to ensure it's a dictionary and contains the expected structure.

Add validation at the beginning of the function:

 def trainer_mapper(legacy_config: dict) -> TrainerConfig:
+    """Convert a legacy configuration dictionary to a TrainerConfig instance.
+
+    This function extracts values from a legacy SLEAP configuration dictionary
+    and maps them to the appropriate fields in the TrainerConfig structure.
+
+    Args:
+        legacy_config: A dictionary containing legacy configuration parameters.
+
+    Returns:
+        A TrainerConfig instance populated with values from the legacy config.
+
+    Raises:
+        TypeError: If legacy_config is not a dictionary.
+    """
+    if not isinstance(legacy_config, dict):
+        raise TypeError("legacy_config must be a dictionary")
+
     return TrainerConfig(

---

236-366: Consider using dictionary access helpers to simplify deeply nested dictionary accesses.

The function contains many instances of deeply nested dictionary access using chained .get() calls, which makes the code hard to read. Consider adding a helper function to simplify these accesses.

Add a helper function at the beginning of the file:

def get_nested_dict_value(dictionary: dict, keys: list, default=None):
    """Safely get a value from a nested dictionary.
    
    Args:
        dictionary: The dictionary to search in.
        keys: A list of keys representing the path to the value.
        default: The default value to return if the key path doesn't exist.
        
    Returns:
        The value at the key path, or the default value if not found.
    """
    current = dictionary
    for key in keys:
        if not isinstance(current, dict):
            return default
        current = current.get(key, {})
    return current or default

Then use it in the function:

-            batch_size=legacy_config.get("optimization", {}).get("batch_size", 1),
+            batch_size=get_nested_dict_value(legacy_config, ["optimization", "batch_size"], 1),

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_trainer_config.py (1)

239-297: Consider verifying the val_data_loader configuration as well.

The new test comprehensively checks most mapped parameters, but it might be beneficial to also verify that val_data_loader is assigned the expected defaults (e.g., batch_size, shuffle, etc.) since it mirrors the training data loader settings. This can increase coverage and potentially catch future regressions.

sleap_nn/config/trainer_config.py (1)

234-367: Remove or implement the commented-out fields for clarity.

Most fields are commented out (e.g., num_workers, amsgrad, step_lr) and can cause confusion about whether they are fully supported. Consider either removing them if they are no longer needed or fully implementing them. This will help maintain code clarity and consistency with the rest of the configuration.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_trainer_config.py (1)

239-297: Consider expanding test coverage for other checkpoint settings.
This test function looks comprehensive, verifying both mapped values and defaults. As an enhancement, consider adding an additional test scenario where "checkpointing" contains "latest_model": True" to confirm config.save_ckpt is correctly set to True.

+ legacy_config["optimization"]["checkpointing"] = {"latest_model": True}
+ config = trainer_mapper(legacy_config)
+ assert config.save_ckpt is True

tests/config/test_model_config.py (1)

123-171: Add optional scenarios for other heads.
The current test thoroughly verifies the single_instance head when using unet. A more exhaustive approach might also test other heads or backbone types if legacy configurations include them (e.g., centroid, centered_instance, multi_instance).

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (2)

277-284: 🛠️ Refactor suggestion

Fix scale parameter to use valid values instead of None.

The scale parameter is set with potentially None values, which could cause validation errors since the GeometricConfig class expects a list of floats with default (0.9, 1.1).

                    scale=(
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_min", None),
+                        .get("scale_min", 0.9),
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_max", None),
+                        .get("scale_max", 1.1),
                    ),

---

297-300: 🛠️ Refactor suggestion

Fix inconsistent handling of use_augmentations_train.

The function:

1. Has a commented-out conditional check for augmentation_config (lines 297-298)
2. Hardcodes use_augmentations_train=True (line 300)

This creates inconsistency with the DataConfig class default (False) and makes the augmentation configuration ignore the actual setting.

            )
-            # if legacy_config.get("use_augmentations_train", False)
-            # else None
+            if legacy_config.get("optimization", {}).get("use_augmentations", True)
+            else None
        ),
-        use_augmentations_train=True,
+        use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),

tests/assets/fixtures/datasets.py (1)

4-4: Remove the unused import to adhere to cleanliness.

Line 4 imports OmegaConf but it's never used in this file. Please remove it to keep the file lean and avoid unnecessary dependencies.

- from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

4-4: omegaconf.OmegaConf imported but unused

Remove unused import: omegaconf.OmegaConf

(F401)

tests/config/test_training_job_config.py (3)

36-36: Remove the unused import for sleapnn_data_dir if it's not needed.

Static analysis indicates sleapnn_data_dir is unused. Review whether this fixture is needed; if not, remove it to eliminate the lint warning and keep dependencies minimal.

- from tests.assets.fixtures.datasets import sleapnn_data_dir, training_job_config_path
+ from tests.assets.fixtures.datasets import training_job_config_path

🧰 Tools

🪛 Ruff (0.8.2)

36-36: tests.assets.fixtures.datasets.sleapnn_data_dir imported but unused

Remove unused import

(F401)

---

36-36: tests.assets.fixtures.datasets.training_job_config_path imported but unused

Remove unused import

(F401)

---

38-38: Remove the unused import dataclasses.asdict.

This import does not appear to be used in the test file. Removing it helps tidy up the code.

- from dataclasses import asdict

🧰 Tools

🪛 Ruff (0.8.2)

38-38: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

41-41: Remove the unused import json.

Although load_sleap_config relies on json, this test file itself does not call json directly. Removing the local import is recommended to satisfy static analysis warnings.

- import json

🧰 Tools

🪛 Ruff (0.8.2)

41-41: json imported but unused

Remove unused import: json

(F401)

sleap_nn/config/model_config.py (1)

837-1003: Consider broader support for non-UNet backbones.

Currently, this mapper only populates a UNetConfig and skips other potential backbone structures. If your legacy configs might include convnext or swint entries, consider handling them similarly. This ensures graceful fallback rather than silently ignoring those entries.

sleap_nn/config/trainer_config.py (3)

237-371: Add a docstring to the trainer_mapper function.

Adding a docstring would improve the maintainability and readability of this function by documenting its purpose, parameters, and return value.

+"""
+Maps a legacy configuration dictionary to a TrainerConfig object.
+
+Parameters:
+    legacy_config (dict): A dictionary containing legacy configuration settings.
+
+Returns:
+    TrainerConfig: An instance of TrainerConfig populated with values from legacy_config.
+"""
def trainer_mapper(legacy_config: dict) -> TrainerConfig:

---

305-309: Simplify the regex for capitalizing the first letter.

The regex to capitalize the first letter of the optimizer name is more complex than needed. Using the built-in capitalize() method would be more readable.

-        optimizer_name=re.sub(
-            r"^[a-z]",
-            lambda x: x.group().upper(),
-            legacy_config.get("optimization", {}).get("optimizer", "adam"),
-        ),
+        optimizer_name=legacy_config.get("optimization", {}).get("optimizer", "adam").capitalize(),

---

242-247: Explain the purpose of commented-out code sections.

There are multiple sections of commented-out code without explanation. This makes it difficult to understand whether these sections are incomplete implementations, deprecated code, or placeholders for future development.

Add explanatory comments at the beginning of commented-out sections, for example:

+        # NOTE: The following parameters are intentionally not mapped as they are not
+        # present in legacy SLEAP configurations or have different default values
         # num_workers=legacy_config.get("optimization", {})
         # .get("train_data_loader", {})
         # .get("num_workers", 0),

Also applies to: 250-255, 258-260, 262-268, 271-281, 282-304

sleap_nn/config/data_config.py (2)

190-191: Add a docstring to the data_mapper function.

The function lacks documentation explaining its purpose, parameters, and return value, which would improve maintainability.

+"""
+Maps a legacy configuration dictionary to a DataConfig object.
+
+Parameters:
+    legacy_config (dict): A dictionary containing legacy configuration settings.
+
+Returns:
+    DataConfig: An instance of DataConfig populated with values from legacy_config.
+"""
def data_mapper(legacy_config: dict) -> DataConfig:

---

236-239: Document the normalization of uniform_noise_max_val.

The function divides uniform_noise_max_val by 100.0, but this normalization is not documented and isn't applied to other similar values.

                    uniform_noise_max=legacy_config.get("optimization", {})
                    .get("augmentation_config", {})
-                    .get("uniform_noise_max_val", 100.0)
-                    / 100.0,
+                    .get("uniform_noise_max_val", 100.0) / 100.0,  # Convert percentage to fraction

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_training_job_config.py (1)

236-279: 🛠️ Refactor suggestion

Improve exception handling and file management in test.

The test has a few issues that should be addressed:

1. Unused variable e in the exception handler
2. Unsafe temporary file handling that doesn't guarantee cleanup
3. Complex control flow with try-except block

Consider using a context manager for temporary file handling and simplifying the test:

def test_load_sleap_config_from_file(training_job_config_path):
    """Test the load_sleap_config function with a sample legacy configuration from a JSON file."""
    # Path to the training_config.json file
    json_file_path = training_job_config_path

    # Load the configuration using the load_sleap_config method
    try:
        # Load the configuration using the load_sleap_config method
        config = load_sleap_config(TrainingJobConfig, json_file_path)
-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)

        # Create a temporary file to hold the modified configuration
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.json', mode='w') as temp_file:
+        with tempfile.NamedTemporaryFile(suffix='.json', mode='w', delete=False) as temp_file:
            old_config['data']['labels']['training_labels'] = "notMISSING"
            old_config['data']['labels']['validation_labels'] = "notMISSING"
            
            json.dump(old_config, temp_file)
            temp_file_path = temp_file.name

-        config = load_sleap_config(TrainingJobConfig, temp_file_path)
-        os.remove(temp_file_path)
+        try:
+            config = load_sleap_config(TrainingJobConfig, temp_file_path)
+        finally:
+            os.remove(temp_file_path)

Also consider validating the exception message to ensure the correct failure is caught:

def test_load_sleap_config_missing_fields(training_job_config_path):
    """Test that load_sleap_config raises appropriate exceptions for missing fields."""
    with tempfile.NamedTemporaryFile(suffix='.json', mode='w') as temp_file:
        # Create a config with missing mandatory fields
        json.dump({"data": {"labels": {}}}, temp_file)
        temp_file.flush()
        
        # Verify the exception is raised with the expected message
        with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
            load_sleap_config(TrainingJobConfig, temp_file.name)

🧰 Tools

🪛 Ruff (0.8.2)

236-236: Redefinition of unused training_job_config_path from line 36

(F811)

---

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

tests/config/test_data_config.py (1)

127-200: Good test coverage for data_mapper function!

The test provides comprehensive validation for all aspects of the data_mapper function, including preprocessing, augmentation, and skeleton configurations.

Consider adding tests for error conditions:

1. What happens when mandatory fields like training_labels or validation_labels are missing?
2. How does the function handle edge cases or invalid values?

def test_data_mapper_missing_mandatory_fields():
    """Test the data_mapper function with missing mandatory fields."""
    # Test missing training_labels
    legacy_config = {
        "data": {
            "labels": {
                "validation_labels": "valid_path",
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
        data_mapper(legacy_config)
    
    # Test missing validation_labels
    legacy_config = {
        "data": {
            "labels": {
                "training_labels": "valid_path",
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="val_labels_path is missing"):
        data_mapper(legacy_config)

tests/config/test_training_job_config.py (3)

31-36: Clean up unused imports.

Several imports are not being used in this file. Consider removing them to improve code clarity.

from sleap_nn.config.training_job_config import TrainingJobConfig
from sleap_nn.config.training_job_config import load_sleap_config
from sleap_nn.config.model_config import ModelConfig
from sleap_nn.config.data_config import DataConfig
from sleap_nn.config.trainer_config import TrainerConfig, EarlyStoppingConfig
from sleap_nn.config.data_config import IntensityConfig
-from tests.assets.fixtures.datasets import sleapnn_data_dir, training_job_config_path
+from tests.assets.fixtures.datasets import training_job_config_path
from omegaconf import OmegaConf, MissingMandatoryValue
-from dataclasses import asdict
import json
-from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

36-36: tests.assets.fixtures.datasets.sleapnn_data_dir imported but unused

Remove unused import

(F401)

---

36-36: tests.assets.fixtures.datasets.training_job_config_path imported but unused

Remove unused import

(F401)

---

245-245: Remove unused exception variable.

The exception variable e is not being used in the exception handler.

    try:
        # Load the configuration using the load_sleap_config method
        config = load_sleap_config(TrainingJobConfig, json_file_path)
-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

251-260: Use context manager for safer temporary file handling.

The current approach doesn't guarantee the temporary file will be deleted if an exception occurs during loading.

-        with tempfile.NamedTemporaryFile(delete=False, suffix='.json', mode='w') as temp_file:
-            old_config['data']['labels']['training_labels'] = "notMISSING"
-            old_config['data']['labels']['validation_labels'] = "notMISSING"
-            
-            json.dump(old_config, temp_file)
-            temp_file_path = temp_file.name
-
-        config = load_sleap_config(TrainingJobConfig, temp_file_path)
-        os.remove(temp_file_path)
+        with tempfile.NamedTemporaryFile(suffix='.json', mode='w', delete=False) as temp_file:
+            old_config['data']['labels']['training_labels'] = "notMISSING"
+            old_config['data']['labels']['validation_labels'] = "notMISSING"
+            
+            json.dump(old_config, temp_file)
+            temp_file_path = temp_file.name
+
+        try:
+            config = load_sleap_config(TrainingJobConfig, temp_file_path)
+        finally:
+            if os.path.exists(temp_file_path):
+                os.remove(temp_file_path)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/model_config.py (1)

835-836: Remove unnecessary blank lines.

These two blank lines appear to serve no functional purpose. Consider removing them to keep the code more concise and aligned with standard style conventions.

tests/config/test_data_config.py (1)

127-200: Test provides comprehensive coverage of the data_mapper function.

The test effectively validates all aspects of the data_mapper function from sleap_nn/config/data_config.py, including:

1. Correct handling of labels paths
2. Preprocessing configuration conversion
3. Augmentation configuration handling
4. Intensity configuration parameters
5. Geometric transformation parameters
6. Skeletons configuration mapping

Consider adding test cases for error conditions, such as missing mandatory values (e.g., when train_labels_path or val_labels_path are missing). This would help ensure the function correctly raises appropriate exceptions in those scenarios.

def test_data_mapper_missing_mandatory_values():
    """Test that data_mapper raises appropriate exceptions for missing mandatory values."""
    # Test missing train_labels_path
    legacy_config_missing_train = {
        "data": {
            "labels": {
                "validation_labels": "some_path.slp"
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="train_labels_path is missing"):
        data_mapper(legacy_config_missing_train)
    
    # Test missing val_labels_path
    legacy_config_missing_val = {
        "data": {
            "labels": {
                "training_labels": "some_path.slp"
            }
        }
    }
    with pytest.raises(MissingMandatoryValue, match="val_labels_path is missing"):
        data_mapper(legacy_config_missing_val)

tests/assets/fixtures/datasets.py (2)

4-4: Remove unused import.
The static analysis tool detects that OmegaConf is imported but unused.

Apply this diff to remove the unused import:

-from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

4-4: omegaconf.OmegaConf imported but unused

Remove unused import: omegaconf.OmegaConf

(F401)

---

34-35: Update docstring to match file name.
The docstring says "Path to centered_instance_training_config file." but the actual file is "centered_instance_with_scaling_training_config.json".

-    """Path to centered_instance_training_config file."""
+    """Path to centered_instance_with_scaling_training_config file."""

tests/config/test_training_job_config.py (8)

36-36: Avoid wildcard import.
Using import * can mask undefined names.

Here is a diff showing an explicit import:

-from tests.assets.fixtures.datasets import *
+from tests.assets.fixtures.datasets import (
+    sleapnn_data_dir,
+    training_job_config_path,
+    bottomup_training_config_path,
+    centered_instance_training_config_path,
+    centered_instance_with_scaling_training_config_path,
+    centroid_training_config_path,
+    single_instance_training_config_path,
+    topdown_training_config_path,
+)

🧰 Tools

🪛 Ruff (0.8.2)

36-36: from tests.assets.fixtures.datasets import * used; unable to detect undefined names

(F403)

---

38-38: Remove unused import.
dataclasses.asdict is not referenced in this file.

-from dataclasses import asdict

🧰 Tools

🪛 Ruff (0.8.2)

38-38: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

42-42: Remove unused import.
omegaconf.MISSING is never used.

-from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

42-42: omegaconf.MISSING imported but unused

Remove unused import: omegaconf.MISSING

(F401)

---

43-43: Remove unused import.
pprint.pprint is never used.

-from pprint import pprint

🧰 Tools

🪛 Ruff (0.8.2)

43-43: pprint.pprint imported but unused

Remove unused import: pprint.pprint

(F401)

---

245-245: Remove assignment to unused variable.
The local variable e is assigned but never used.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

296-296: Remove assignment to unused variable.
The local variable e is assigned but never used.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

296-296: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

340-340: Remove assignment to unused variable.
Again, e is not utilized.

-    except MissingMandatoryValue as e:
+    except MissingMandatoryValue:

🧰 Tools

🪛 Ruff (0.8.2)

340-340: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

237-285: Refactor repeated code blocks in tests.
Each test repeats the same logic for handling MissingMandatoryValue, rewriting JSON, and loading the updated file.

Consider extracting this into a helper function or fixture to reduce duplication:

def load_or_fix_legacy_config(json_file_path, required_keys=None):
    """
    Attempt to load the config. If MissingMandatoryValue is raised,
    rewrite the config JSON with non-missing placeholders.
    """
    from omegaconf import MissingMandatoryValue
    import tempfile, os, json

    if required_keys is None:
        required_keys = [("data", "labels", "training_labels"), ("data", "labels", "validation_labels")]
    try:
        return load_sleap_config(TrainingJobConfig, json_file_path)
    except MissingMandatoryValue:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)

        for key_path in required_keys:
            cfg = old_config
            for kp in key_path[:-1]:
                cfg = cfg[kp]
            cfg[key_path[-1]] = "notMISSING"

        with tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w") as temp_file:
            json.dump(old_config, temp_file)
            temp_file_path = temp_file.name

        config = load_sleap_config(TrainingJobConfig, temp_file_path)
        os.remove(temp_file_path)
        return config

Also applies to: 287-327, 329-371, 373-407, 409-439, 441-473, 475-510

🧰 Tools

🪛 Ruff (0.8.2)

245-245: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (1)

190-313: 🛠️ Refactor suggestion

Several concerns in the new data_mapper function.

1. Commented-out code: Key parameters like provider, user_instances_only, etc., remain commented out. If they are required, consider uncommenting or removing them altogether to reduce confusion.
2. Forcing use_augmentations_train to True: You currently override this with a hard-coded True (line 310). Most designs expect reading it from the legacy config (e.g., legacy_config.get("use_augmentations_train", False)).
3. Potential (None, None) usage: Passing (None, None) as scale (lines 288–293) can break validations in GeometricConfig. Confirm whether this is intended or provide a default numeric range if the config is missing a scale setting.

Below is a snippet to let legacy_config decide if augmentations are used, and to supply a safer default for scale:

-        use_augmentations_train=True,
+        use_augmentations_train=bool(
+            legacy_config.get("optimization", {})
+                       .get("use_augmentations", False)
+        ),

-                    scale=(
-                        legacy_config.get("optimization", {})
-                        .get("augmentation_config", {})
-                        .get("scale_min", None),
-                        legacy_config.get("optimization", {})
-                        .get("augmentation_config", {})
-                        .get("scale_max", None),
-                    ),
+                    scale=(
+                        legacy_config.get("optimization", {})
+                        .get("augmentation_config", {})
+                        .get("scale_min", 0.9),
+                        legacy_config.get("optimization", {})
+                        .get("augmentation_config", {})
+                        .get("scale_max", 1.1),
+                    ),

sleap_nn/config/training_job_config.py (1)

155-173: 🛠️ Refactor suggestion

Multiple issues with the load_sleap_config function.

1. Lack of error handling: No try/except to handle file I/O errors and JSON parse errors (e.g., FileNotFoundError, json.JSONDecodeError). This can lead to uncaught exceptions at runtime.
2. @classmethod usage: The function uses cls as its first parameter, but it isn’t decorated with @classmethod, causing potential confusion about its intended usage.
3. Missing check_output_strides: Unlike from_yaml and load_yaml, this loader doesn’t invoke TrainingJobConfig.check_output_strides(config), potentially leading to inconsistent stride settings relative to the YAML loading path.
4. Return type mismatch: The function signature says it returns a TrainerConfig, but the actual return object is an OmegaConf instance.

Below is a suggested refactor that includes error handling, calls check_output_strides, and aligns the return type:

+@classmethod
-def load_sleap_config(cls, json_file_path: str) -> TrainerConfig:
+def load_sleap_config(cls, json_file_path: str) -> OmegaConf:
    +import json
    try:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)
    except FileNotFoundError:
        raise FileNotFoundError(f"Configuration file not found: {json_file_path}")
    except json.JSONDecodeError:
        raise ValueError(f"Invalid JSON in configuration file: {json_file_path}")

    data_config = data_mapper(old_config)
    model_config = model_mapper(old_config)
    trainer_config = trainer_mapper(old_config)

    config = cls(
        data_config=data_config,
        model_config=model_config,
        trainer_config=trainer_config,
    )

    schema = OmegaConf.structured(config)
    config_omegaconf = OmegaConf.merge(schema, OmegaConf.create(asdict(config)))
    +config_omegaconf = cls.check_output_strides(config_omegaconf)
    OmegaConf.to_container(config_omegaconf, resolve=True, throw_on_missing=True)

    return config_omegaconf

sleap_nn/config/data_config.py (1)

8-8: Remove the unused import.
MissingMandatoryValue is never invoked, so please remove it to address the linter’s warning and keep your imports clean.

-from omegaconf import MISSING, MissingMandatoryValue
+from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

8-8: omegaconf.MissingMandatoryValue imported but unused

Remove unused import: omegaconf.MissingMandatoryValue

(F401)

tests/config/test_training_job_config.py (2)

36-37: Avoid star-import if possible.
Using import * from fixtures can clutter the namespace and make it harder to track dependencies. Prefer explicitly importing only what you need.

🧰 Tools

🪛 Ruff (0.8.2)

36-36: from tests.assets.fixtures.datasets import * used; unable to detect undefined names

(F403)

---

41-43: Several unused imports.

• json is used in the tests below, so it’s fine.
• MISSING (line 42) and pprint (line 43) appear unused. Consider removing them along with any other imports flagged by static analysis for clarity.

 import json
-from omegaconf import MISSING
-from pprint import pprint

🧰 Tools

🪛 Ruff (0.8.2)

42-42: omegaconf.MISSING imported but unused

Remove unused import: omegaconf.MISSING

(F401)

---

43-43: pprint.pprint imported but unused

Remove unused import: pprint.pprint

(F401)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (2)

295-302: 🛠️ Refactor suggestion

Fix scale parameter default values

The scale parameter in GeometricConfig is set to None values, which could conflict with validators that expect numeric values. In the GeometricConfig class definition (line 113), the default is (0.9, 1.1).

                    scale=(
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_min", None),
+                        .get("scale_min", 0.9),
                        legacy_config.get("optimization", {})
                        .get("augmentation_config", {})
-                        .get("scale_max", None),
+                        .get("scale_max", 1.1),
                    ),

---

315-318: ⚠️ Potential issue

Fix inconsistent use_augmentations_train handling

There are two issues:

1. use_augmentations_train is hardcoded to True regardless of what's in the legacy configuration
2. The conditional check for setting augmentation_config to None is commented out

This could lead to inconsistencies between the augmentation configuration and whether augmentations are actually used.

            )
-            # if legacy_config.get("use_augmentations_train", False)
-            # else None
+            if legacy_config.get("optimization", {}).get("use_augmentations", True)
+            else None
        ),
-        use_augmentations_train=True,
+        use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),

sleap_nn/config/trainer_config.py (4)

313-317: Consider simplifying the regex for capitalizing the optimizer name.

The regex approach works but could be simplified using standard string methods.

-    optimizer_name=re.sub(
-        r"^[a-z]",
-        lambda x: x.group().upper(),
-        legacy_config.get("optimization", {}).get("optimizer", "adam"),
-    ),
+    optimizer_name=legacy_config.get("optimization", {}).get("optimizer", "adam").capitalize(),

---

270-311: Consider adding a TODO comment for commented-out configurations.

There are numerous commented-out configurations in this function. If these are planned for future implementation, consider adding a TODO comment explaining the plan. If they're not needed, consider removing them for better code clarity.

---

324-363: Consider adding a helper function for nested dictionary access.

The deeply nested get() calls reduce readability. Consider extracting a helper function to make the code more maintainable.

def get_nested(d, *keys, default=None):
    """Get a value from a nested dictionary by a sequence of keys."""
    result = d
    for key in keys:
        if not isinstance(result, dict):
            return default
        result = result.get(key, {})
    return result if result != {} else default

# Usage example:
patience = get_nested(legacy_config, "optimization", "learning_rate_schedule", "plateau_patience", default=10)

---

349-358: Inconsistent path access could lead to confusion.

Some parameters are accessed from "learning_rate_schedule" while others might come from "lr_scheduler" (in commented code). This inconsistency could lead to confusion.

Consider documenting why different paths are used or standardize on a consistent approach if possible.

sleap_nn/config/model_config.py (4)

846-895: Consider simplifying repetitive dictionary access patterns

The code contains many repetitive nested .get() calls which make the code verbose and harder to maintain. Consider creating a helper function to simplify this pattern.

+def _get_nested(config, path, default=None):
+    """Get a value from a nested dictionary using a list of keys.
+    
+    Args:
+        config: Dictionary to access
+        path: List of keys to traverse
+        default: Default value if path doesn't exist
+        
+    Returns:
+        Value at the path or default if not found
+    """
+    result = config
+    for key in path:
+        if not isinstance(result, dict):
+            return default
+        result = result.get(key, {})
+    return result if result != {} else default

 def model_mapper(legacy_config: dict) -> ModelConfig:
     """Map the legacy model configuration to the new model configuration.
 
     Args:
         legacy_config: A dictionary containing the legacy model configuration.
 
     Returns:
         An instance of `ModelConfig` with the mapped configuration.
     """
+    model_config = legacy_config.get("model", {})
     return ModelConfig(
         # init_weights=legacy_config.get("init_weights", "default"),
         # pre_trained_weights not in old config
         # pretrained_backbone_weights=legacy_config.get("PretrainedEncoderConfig")?? # i think its different
         # pretrained_head_weights not in old config
         backbone_config=BackboneConfig(
             unet=(
                 UNetConfig(
                     # in_channels=legacy_config.get("backbone", {}).get("in_channels", 1),
                     # kernel_size=legacy_config.get("backbone", {}).get("kernel_size", 3),
-                    filters=legacy_config.get("model", {})
-                    .get("backbone", {})
-                    .get("unet", {})
-                    .get("filters", 32),
+                    filters=_get_nested(model_config, ["backbone", "unet", "filters"], 32),

This pattern could be applied throughout the function to make it more maintainable.

---

896-920: Remove unnecessary nested parentheses

The SingleInstanceConfig has an unnecessary level of parentheses in the conditional, which can be simplified.

         head_configs=HeadConfig(
             single_instance=(
-                (
-                    SingleInstanceConfig(
-                        confmaps=SingleInstanceConfMapsConfig(
-                            part_names=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("part_names"),
-                            sigma=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("sigma", 5.0),
-                            output_stride=legacy_config.get("model", {})
-                            .get("heads", {})
-                            .get("single_instance", {})
-                            .get("output_stride", 1),
-                        )
-                    )
-                )
+                SingleInstanceConfig(
+                    confmaps=SingleInstanceConfMapsConfig(
+                        part_names=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("part_names"),
+                        sigma=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("sigma", 5.0),
+                        output_stride=legacy_config.get("model", {})
+                        .get("heads", {})
+                        .get("single_instance", {})
+                        .get("output_stride", 1),
+                    )
+                )
                 if legacy_config.get("model", {})
                 .get("heads", {})
                 .get("single_instance")
                 else None
             ),

---

969-1016: The bottomup head mapping uses multi_instance paths

Good job on adding the conditional check to maintain the @oneof constraint for the bottomup head. I noticed that the legacy configuration key is multi_instance while the new configuration key is bottomup. This seems intentional, but it would be helpful to add a comment explaining this mapping for future maintenance.

             bottomup=(
                 BottomUpConfig(
+                    # Note: bottomup head was previously called multi_instance in legacy configs
                     confmaps=BottomUpConfMapsConfig(
                         loss_weight=legacy_config.get("model", {})
                         .get("heads", {})
                         .get("multi_instance", {})

---

837-1018: Add input validation for legacy_config

The function assumes that legacy_config is a valid dictionary, but doesn't validate this. Consider adding input validation to handle edge cases gracefully.

 def model_mapper(legacy_config: dict) -> ModelConfig:
     """Map the legacy model configuration to the new model configuration.
 
     Args:
         legacy_config: A dictionary containing the legacy model configuration.
 
     Returns:
         An instance of `ModelConfig` with the mapped configuration.
     """
+    if not isinstance(legacy_config, dict):
+        logger.warning("Invalid legacy_config provided, using defaults")
+        legacy_config = {}
+    
     return ModelConfig(

sleap_nn/config/data_config.py (4)

8-8: Remove unused import

The MissingMandatoryValue exception is imported but never used in the code.

-from omegaconf import MISSING, MissingMandatoryValue
+from omegaconf import MISSING

🧰 Tools

🪛 Ruff (0.8.2)

8-8: omegaconf.MissingMandatoryValue imported but unused

Remove unused import: omegaconf.MissingMandatoryValue

(F401)

---

209-216: Add explanation for commented-out parameters

Several parameters are commented out without explanation. If these parameters aren't used in legacy configurations, add a comment explaining why they're commented out. If they should be included, uncomment them.

-        # provider=legacy_config.get("provider", "LabelsReader"),
-        # user_instances_only=legacy_config.get("user_instances_only", True),
-        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
-        # np_chunks_path=legacy_config.get("np_chunks_path"),
-        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
-        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
-        # chunk_size=int(legacy_config.get("chunk_size", 100)),
-        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),
+        # These parameters are not part of the legacy SLEAP configuration format
+        # provider=legacy_config.get("provider", "LabelsReader"),
+        # user_instances_only=legacy_config.get("user_instances_only", True),
+        # data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
+        # np_chunks_path=legacy_config.get("np_chunks_path"),
+        # litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
+        # use_existing_chunks=legacy_config.get("use_existing_chunks", False),
+        # chunk_size=int(legacy_config.get("chunk_size", 100)),
+        # delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),

---

303-312: Explain commented-out geometric augmentation parameters

Several geometric augmentation parameters are commented out without explanation. If these parameters aren't available in legacy configurations, add a comment explaining why they're commented out. If they should be included, uncomment them.

-                    # translate_width=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_width", 0.2),
-                    # translate_height=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_height", 0.2),
-                    # affine_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("affine_p", 0.0),
-                    # erase_scale_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_min", 0.0001),
-                    # erase_scale_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_max", 0.01),
-                    # erase_ratio_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_min", 1.0),
-                    # erase_ratio_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_max", 1.0),
-                    # erase_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_p", 0.0),
-                    # mixup_lambda=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_lambda", [0.01, 0.05]),
-                    # mixup_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_p", 0.0),
+                    # These geometric augmentation parameters are not available in legacy SLEAP configuration format
+                    # translate_width=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_width", 0.2),
+                    # translate_height=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("translate_height", 0.2),
+                    # affine_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("affine_p", 0.0),
+                    # erase_scale_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_min", 0.0001),
+                    # erase_scale_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_scale_max", 0.01),
+                    # erase_ratio_min=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_min", 1.0),
+                    # erase_ratio_max=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_ratio_max", 1.0),
+                    # erase_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("erase_p", 0.0),
+                    # mixup_lambda=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_lambda", [0.01, 0.05]),
+                    # mixup_p=legacy_config.get("optimization", {}).get("augmentation_config", {}).get("mixup_p", 0.0),

---

318-318: Add unit tests for data_mapper function

The data_mapper function is complex and handles critical configuration mapping, but there don't appear to be tests for it. Consider adding unit tests to verify it handles various legacy configuration formats correctly.

Would you like me to generate example unit tests for the data_mapper function?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/data_config.py (3)

187-187: ⚠️ Potential issue

Type annotation mismatch for skeletons.

The skeletons field is typed as Optional[dict] here, but the function implementation on line 321 extracts it from a list. Based on the test_data_mapper test, the expected behavior is to return a dictionary.

---

296-302: ⚠️ Potential issue

Invalid default values for scale parameter.

The scale parameter is using None as potential default values, which might conflict with the validators in the GeometricConfig class. The test expects specific float values.

Apply this diff to fix the scale parameter:

scale=(
-    legacy_config.get("optimization", {})
-    .get("augmentation_config", {})
-    .get("scale_min", None),
-    legacy_config.get("optimization", {})
-    .get("augmentation_config", {})
-    .get("scale_max", None),
+    legacy_config.get("optimization", {})
+    .get("augmentation_config", {})
+    .get("scale_min", 0.9),
+    legacy_config.get("optimization", {})
+    .get("augmentation_config", {})
+    .get("scale_max", 1.1),
),

---

237-237: 🛠️ Refactor suggestion

Inconsistent handling of use_augmentations_train.

The use_augmentations_train parameter is commented out here, but then hardcoded to True on line 318. This seems inconsistent and might lead to unexpected behavior.

Either uncomment this line to use the value from the legacy config or add a comment explaining why it's hardcoded:

-# use_augmentations_train=legacy_config.get("use_augmentations_train", False),
+use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),

sleap_nn/config/data_config.py (2)

209-216: Commented out parameters without explanation.

Multiple important parameters are commented out without explanation. If these parameters should be excluded from the legacy configuration mapping, consider adding a comment explaining why.

Either uncomment these parameters if they should be included, or add a comment explaining why they're excluded:

+# The following parameters are intentionally excluded as they are not present in legacy configs
# provider=legacy_config.get("provider", "LabelsReader"),
# user_instances_only=legacy_config.get("user_instances_only", True),
# data_pipeline_fw=legacy_config.get("data_pipeline_fw", "torch_dataset"),
# np_chunks_path=legacy_config.get("np_chunks_path"),
# litdata_chunks_path=legacy_config.get("litdata_chunks_path"),
# use_existing_chunks=legacy_config.get("use_existing_chunks", False),
# chunk_size=int(legacy_config.get("chunk_size", 100)),
# delete_chunks_after_training=legacy_config.get("delete_chunks_after_training", True),

---

315-317: Commented out conditional logic for augmentation_config.

The conditional logic for augmentation_config is commented out, which means it will always be created regardless of the value of use_augmentations_train in the legacy configuration.

Either uncomment this conditional logic or remove it entirely if augmentation_config should always be created:

)
-# if legacy_config.get("use_augmentations_train", False)
-# else None
+if legacy_config.get("optimization", {}).get("use_augmentations", True)
+else None
),

tests/config/test_data_config.py (1)

127-200: Good test coverage but missing edge cases.

The test function provides good coverage for the happy path, but doesn't test any error cases. Consider adding tests for:

1. Missing required fields (train_labels_path, val_labels_path)
2. Different skeletons formats (dict vs list)
3. Invalid values that might trigger validators

Add additional test cases to cover error scenarios:

def test_data_mapper_missing_required_fields():
    """Test the data_mapper function with missing required fields."""
    legacy_config = {
        "data": {
            "labels": {
                # Missing training_labels and validation_labels
                "skeletons": [{"edges": [[0, 1], [1, 2]]}],
            }
        }
    }
    
    with pytest.raises(MissingMandatoryValue):
        data_mapper(legacy_config)

def test_data_mapper_different_skeleton_formats():
    """Test the data_mapper function with different skeleton formats."""
    # Test with skeletons as a dictionary
    legacy_config_dict = {
        "data": {
            "labels": {
                "training_labels": "notMISSING",
                "validation_labels": "notMISSING",
                "skeletons": {"edges": [[0, 1], [1, 2]]},
            }
        }
    }
    
    config_dict = data_mapper(legacy_config_dict)
    assert config_dict.skeletons == {"edges": [[0, 1], [1, 2]]}
    
    # Test with skeletons as a list
    legacy_config_list = {
        "data": {
            "labels": {
                "training_labels": "notMISSING",
                "validation_labels": "notMISSING",
                "skeletons": [{"edges": [[0, 1], [1, 2]]}],
            }
        }
    }
    
    config_list = data_mapper(legacy_config_list)
    assert config_list.skeletons == {"edges": [[0, 1], [1, 2]]}

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/training_job_config.py (2)

155-182: ⚠️ Potential issue

Missing validation and error handling in load_sleap_config_from_json

The function needs to add error handling for file operations and JSON parsing and apply the check_output_strides validation after configuration loading.

Apply this diff to fix both issues:

def load_sleap_config_from_json(cls, json_file_path: str) -> OmegaConf:
    """Load a SLEAP configuration from a JSON file and convert it to OmegaConf.

    Args:
        cls: The class to instantiate with the loaded configuration.
        json_file_path: Path to a JSON file containing the SLEAP configuration.

    Returns:
        An OmegaConf instance with the loaded configuration.
    """
-    with open(json_file_path, "r") as f:
-        old_config = json.load(f)
+    try:
+        with open(json_file_path, "r") as f:
+            old_config = json.load(f)
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Configuration file not found: {json_file_path}")
+    except json.JSONDecodeError:
+        raise ValueError(f"Invalid JSON in configuration file: {json_file_path}")

    data_config = data_mapper(old_config)
    model_config = model_mapper(old_config)
    trainer_config = trainer_mapper(old_config)

    config = cls(
        data_config=data_config,
        model_config=model_config,
        trainer_config=trainer_config,
    )

    schema = OmegaConf.structured(config)
    config_omegaconf = OmegaConf.merge(schema, OmegaConf.create(asdict(config)))
    OmegaConf.to_container(config_omegaconf, resolve=True, throw_on_missing=True)
+    config_omegaconf = TrainingJobConfig.check_output_strides(config_omegaconf)

    return config_omegaconf

---

155-182: 🛠️ Refactor suggestion

Function signature issues in load_sleap_config_from_json

The function seems intended to be a class method but lacks the @classmethod decorator, and the return type annotation should be more specific.

Consider moving this function into the TrainingJobConfig class and fixing the return type:

-def load_sleap_config_from_json(cls, json_file_path: str) -> OmegaConf:
+@classmethod
+def load_sleap_config_from_json(cls, json_file_path: str) -> "TrainingJobConfig":

sleap_nn/config/data_config.py (2)

318-321: ⚠️ Potential issue

Hardcoded use_augmentations_train and inconsistent skeletons extraction

Two issues in this section:

1. use_augmentations_train is hardcoded to True instead of getting it from the legacy configuration
2. The skeletons extraction assumes it's a list and takes the first item, which might cause issues if the structure is different

Apply this diff to fix the issues:

-use_augmentations_train=True,
+use_augmentations_train=legacy_config.get("optimization", {}).get("use_augmentations", True),
skeletons=legacy_config.get("data", {})
.get("labels", {})
-.get("skeletons", [{}])[0],
+.get("skeletons", {}) if isinstance(legacy_config.get("data", {}).get("labels", {}).get("skeletons"), dict) else 
+legacy_config.get("data", {}).get("labels", {}).get("skeletons", [{}])[0],

---

295-302: 🛠️ Refactor suggestion

Fix potential validator issue with GeometricConfig scale parameter

The scale parameter in GeometricConfig uses None as potential values which might conflict with the validator.

Apply this diff to fix the issue:

scale=(
    legacy_config.get("optimization", {})
    .get("augmentation_config", {})
-    .get("scale_min", None),
+    .get("scale_min", 0.9),
    legacy_config.get("optimization", {})
    .get("augmentation_config", {})
-    .get("scale_max", None),
+    .get("scale_max", 1.1),
),

sleap_nn/config/training_job_config.py (1)

155-182: Consider renaming the function for consistency

The function name load_sleap_config_from_json is inconsistent with other similar methods like load_yaml. Consider a more consistent naming pattern.

-def load_sleap_config_from_json(cls, json_file_path: str) -> OmegaConf:
+def load_json(cls, json_file_path: str) -> OmegaConf:

sleap_nn/config/data_config.py (1)

209-216: Explain commented-out code in data_mapper function

There are multiple sections of commented-out code in the data_mapper function without explanation. This includes configuration parameters (lines 209-216), conditional code (lines 315-316), and the use_augmentations_train mapping (line 237).

Consider either:

1. Removing the commented code if it's not needed
2. Uncommenting the code if it should be functional
3. Adding explanatory comments about why this code is commented out

For instance, if these parameters aren't mapped from legacy configuration:

# These parameters are not available in legacy SLEAP configs and use defaults
# provider=legacy_config.get("provider", "LabelsReader"),
# user_instances_only=legacy_config.get("user_instances_only", True),

Also applies to: 237-238, 315-316

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/config/test_training_job_config.py (1)

246-265: 🛠️ Refactor suggestion

Refactor repeated validation error handling pattern.

This pattern of catching ValidationError, modifying the config, and retrying is repeated in all seven test functions. Consider extracting it to a helper function.

The exception variable e is caught but never used:

-try:
-    config = load_sleap_config(TrainingJobConfig, json_file_path)
-except ValidationError as e:
+try:
+    config = load_sleap_config(TrainingJobConfig, json_file_path)
+except ValidationError:

Consider implementing a helper function like:

def load_config_with_fallback(json_file_path):
    """Load config and handle missing mandatory fields."""
    try:
        return load_sleap_config(TrainingJobConfig, json_file_path)
    except ValidationError:
        with open(json_file_path, "r") as f:
            old_config = json.load(f)

        with tempfile.NamedTemporaryFile(
            delete=False, suffix=".json", mode="w"
        ) as temp_file:
            old_config["data"]["labels"]["training_labels"] = "notMISSING"
            old_config["data"]["labels"]["validation_labels"] = "notMISSING"
            json.dump(old_config, temp_file)
            temp_file_path = temp_file.name

        config = load_sleap_config(TrainingJobConfig, temp_file_path)
        os.remove(temp_file_path)
        return config

This would reduce duplication significantly.

🧰 Tools

🪛 Ruff (0.8.2)

248-248: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

tests/assets/fixtures/datasets.py (1)

35-36: Inconsistent docstring for centered_instance_with_scaling_training_config_path.

The docstring doesn't match the actual functionality - it references "centered_instance_training_config" instead of "centered_instance_with_scaling_training_config".

-    """Path to centered_instance_training_config file."""
+    """Path to centered_instance_with_scaling_training_config file."""

tests/config/test_training_job_config.py (3)

36-43: Clean up imports.

Several imports are unused and should be removed:

1. The wildcard import could be replaced with specific imports
2. Remove unused imports: asdict, MISSING, and pprint

-from tests.assets.fixtures.datasets import *
-from omegaconf import OmegaConf, MissingMandatoryValue, ValidationError
-from dataclasses import asdict
+from tests.assets.fixtures.datasets import (
+    bottomup_multiclass_training_config_path,
+    bottomup_training_config_path,
+    centered_instance_training_config_path,
+    centered_instance_with_scaling_training_config_path,
+    centroid_training_config_path,
+    single_instance_training_config_path,
+    topdown_training_config_path,
+)
+from omegaconf import OmegaConf, MissingMandatoryValue, ValidationError
 from loguru import logger
 from _pytest.logging import LogCaptureFixture
 import json
-from omegaconf import MISSING
-from pprint import pprint

🧰 Tools

🪛 Ruff (0.8.2)

36-36: from tests.assets.fixtures.datasets import * used; unable to detect undefined names

(F403)

---

38-38: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

42-42: omegaconf.MISSING imported but unused

Remove unused import: omegaconf.MISSING

(F401)

---

43-43: pprint.pprint imported but unused

Remove unused import: pprint.pprint

(F401)

---

504-511: Uncomment or remove commented test assertions.

The commented assertion lines for the topdown head configuration suggest incomplete tests. Either uncomment and complete these assertions or remove them if they're not needed.

---

237-512: Tests are comprehensive but have redundancy.

The tests cover multiple configuration types which is good for comprehensive testing. However, there's significant repetition in the test structure. Consider creating parameterized tests to reduce duplication.

You could use pytest's parameterization to reduce the repetitive test code:

@pytest.mark.parametrize(
    "config_path_fixture,expected_values",
    [
        ("bottomup_multiclass_training_config_path", {
            "backbone.unet.filters": 8,
            "backbone.unet.max_stride": 16,
            # other expected values...
        }),
        ("bottomup_training_config_path", {
            "backbone.unet.filters": 16,
            "backbone.unet.max_stride": 8,
            # other expected values...
        }),
        # other configurations...
    ],
)
def test_load_training_config_from_file(config_path_fixture, expected_values, request):
    """Test loading configurations from JSON files."""
    json_file_path = request.getfixturevalue(config_path_fixture)
    
    # Use the helper function mentioned earlier
    config = load_config_with_fallback(json_file_path)
    
    # Assert expected values
    for path, value in expected_values.items():
        assert _get_nested_attribute(config, path) == value
        
def _get_nested_attribute(obj, path):
    """Access nested attributes using dot notation."""
    parts = path.split(".")
    for part in parts:
        obj = getattr(obj, part)
    return obj

🧰 Tools

🪛 Ruff (0.8.2)

248-248: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

299-299: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

---

343-343: Local variable e is assigned to but never used

Remove assignment to unused variable e

(F841)

tests/assets/bottomup_multiclass_training_config.json (3)

1-80: Comprehensive "data" Section – Verify Legacy Values and Types

The "data" section is well structured with clear subdivisions for labels, skeletons, and preprocessing parameters. A couple of points to consider:

• Verify that the use of null for "training_labels", "validation_labels", and "test_labels" is intentional. In some cases, an empty list may be preferable to avoid downstream deserialization issues.
• The legacy serialization details (e.g., "py/object", "py/state", "py/tuple") in the skeletons section should be properly mapped to the new config schema. Confirm that your mapping function converts these fields as needed by the new sleapnn requirements.

---

81-125: Robust "model" Section – Confirm Handling of Null and Detailed Configurations

The "model" portion defines detailed configurations for backbones and heads (including the multi-class bottom-up setup). Please ensure that:

• All keys explicitly set to null (for example, "leap", "hourglass", "resnet", "pretrained_encoder") are correctly interpreted by your mapping function.
• The nested configuration under "multi_class_bottomup" (with "confmaps" and "class_maps") is fully supported in the new configuration object.

---

126-185: Thorough "Optimization" Section – Validate Parameter Formats and Defaults

The "optimization" section is comprehensive and includes augmentation settings, learning rate scheduling, and more. A couple of suggestions:

• Confirm that numerical values in scientific notation (e.g., 1e-07) are parsed correctly by the JSON decoder and used appropriately in your new configuration.
• Double-check that default values and parameter ranges (such as the rotation angles, crop sizes, and batch sizes) match the expectations of the new training configurations.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro